# load data
from llama_index.core import SimpleDirectoryReader,VectorStoreIndex, Settings
from llama_index.core.tools.query_engine import QueryEngineTool
from llama_index.core.query_engine import SubQuestionQueryEngine, RouterQueryEngine, ToolRetrieverRouterQueryEngine
from llama_index.core.tools.tool_spec.base import ToolMetadata
from llms import deepseek_llm
from embeddings import embed_model_local_bge_small
Settings.llm = deepseek_llm()
Settings.embed_model= embed_model_local_bge_small()
cities = ["changchun", "hangzhou", "lanzhou", "shenzhen"]
cities_zh = ["长春", "杭州", "兰州", "深圳"]
cities_data = SimpleDirectoryReader(input_dir="./data/").load_data()

# build index and query engine
vector_query_engines = []
for city_data in cities_data:
    vector_query_engines.append(
        # 创建索引（内存）
        VectorStoreIndex.from_documents(
            [city_data],
            use_async=True,
        ).as_query_engine()
    )



# setup base query engine as tool
query_engine_tools =[]
for vector_query_engine, city in zip(vector_query_engines, cities_zh):#vector_query_engines:
    query_engine_tools.append(
        QueryEngineTool(
            query_engine=vector_query_engine,
            metadata=ToolMetadata(
                name="city_data",
                description=f"回答关于 {city} 的任何问题",
            ),
        )
    )

query_engine = RouterQueryEngine.from_defaults(
    query_engine_tools=query_engine_tools,
    verbose=True,
    select_multi=True
)

# 分割成两个问题，并且每个问题都使用不同的工具，然后再合并
response = query_engine.query(
    "介绍一下深圳和杭州的历史"
)
print(response)